package com.zhb.zhbaiagent.rag;

import jakarta.annotation.Resource;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.document.Document;
import org.springframework.ai.transformer.KeywordMetadataEnricher;
import org.springframework.stereotype.Component;

import java.util.List;

@Component
public class MyKeyEnricher {
    //调用灵机大模型
    @Resource
    private ChatModel dashscopeChatModel;
    /**
     * 对文档列表进行关键词元数据增强处理。
     * 使用指定的ChatModel（如灵机大模型）来生成关键词，并将这些关键词作为元数据添加到对应的文档中。
     *
     * @param documents 需要进行关键词增强的原始文档列表
     * @return 经过关键词元数据增强后的文档列表
     */
    public List<Document> enrichDocumets(List<Document> documents) {
        // 创建KeywordMetadataEnricher实例，设置使用的ChatModel和生成关键词的最大数量
        KeywordMetadataEnricher keywordMetadataEnricher = new KeywordMetadataEnricher(dashscopeChatModel, 5);
        // 对输入文档应用关键词增强逻辑并返回结果
        return keywordMetadataEnricher.apply(documents);
    }
}
